module conv #(
    parameter int unsigned P_ICH       = 4,   // Parallel input channels
    parameter int unsigned P_OCH       = 4,   // Parallel output channels
    parameter int unsigned N_ICH       = 16,  // Total input channels
    parameter int unsigned N_OCH       = 16,  // Total output channels
    parameter int unsigned K           = 3,   // Kernel size (KxK)
    parameter int unsigned A_BIT       = 8,   // Input activation bit width
    parameter int unsigned W_BIT       = 8,   // Weight bit width
    parameter int unsigned B_BIT       = 32,  // Output bit width (accumulator)
    parameter int unsigned N_HW        = 64,  // Number of vectors to process
    parameter string       WEIGHT_FILE = ""   // Weight ROM initialization file
) (
    input logic clk,
    input logic rst_n,

    // Input stream interface
    input  logic [P_ICH*A_BIT-1:0] in_data,
    input  logic                   in_valid,
    output logic                   in_ready,

    // Output stream interface
    output logic [P_OCH*B_BIT-1:0] out_data,
    output logic                   out_valid,
    input  logic                   out_ready
);

    // Derived parameters
    localparam int unsigned FOLD_I = N_ICH / P_ICH;
    localparam int unsigned FOLD_O = N_OCH / P_OCH;
    localparam int unsigned KK = K * K;
    localparam int unsigned WEIGHT_DEPTH = FOLD_O * FOLD_I * KK;

    // Weight ROM signals
    logic [$clog2(WEIGHT_DEPTH)-1:0] weight_addr;
    logic [   P_OCH*P_ICH*W_BIT-1:0] weight_data;
    // Weight ROM instantiation
    rom #(
        .DWIDTH(P_OCH * P_ICH * W_BIT),
        .AWIDTH($clog2(WEIGHT_DEPTH)),
        .MEM_SIZE(WEIGHT_DEPTH),
        .INIT_FILE(WEIGHT_FILE)
    ) u_weight_rom (
        .clk  (clk),
        .ce0  (1'b1),
        .addr0(weight_addr),
        .q0   (weight_data)
    );

    // Line buffer to store input data for reuse across output folds
    logic        [     P_ICH*A_BIT-1:0] line_buffer               [ FOLD_I] [KK];

    // Accumulator array
    logic signed [           B_BIT-1:0] acc                       [  P_OCH];

    // Pipeline counters
    logic        [  $clog2(N_HW+1)-1:0] cntr_hw;
    logic        [$clog2(FOLD_O+1)-1:0] cntr_fo;
    logic        [$clog2(FOLD_I+1)-1:0] cntr_fi;
    logic        [    $clog2(KK+1)-1:0] cntr_kk;

    // Pipeline control signals
    logic                               pipeline_enable;
    logic                               is_fst_fo;
    logic                               is_lst_kk;
    logic                               is_lst_fi;

    // Delayed signals for total latency (1 cycle for weight + P_ICH cycles for mac_array)
    logic                               is_lst_kk_dly             [P_ICH+1];
    logic                               is_lst_fi_dly             [P_ICH+1];
    logic                               output_ready_to_write_dly;

    // Assign control signals
    assign is_fst_fo = (cntr_fo == 0);
    assign is_lst_kk = (cntr_kk == KK - 1);
    assign is_lst_fi = (cntr_fi == FOLD_I - 1);

    // Delay chain for is_lst_kk and is_lst_fi to match total latency (1 + P_ICH cycles)
    always_ff @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            for (int i = 0; i <= P_ICH; i++) begin
                is_lst_kk_dly[i] <= 1'b0;
                is_lst_fi_dly[i] <= 1'b0;
            end
        end else if (pipeline_enable) begin
            is_lst_kk_dly[0] <= is_lst_kk;
            is_lst_fi_dly[0] <= is_lst_fi;
            for (int i = 1; i <= P_ICH; i++) begin
                is_lst_kk_dly[i] <= is_lst_kk_dly[i-1];
                is_lst_fi_dly[i] <= is_lst_fi_dly[i-1];
            end
        end
    end

    assign output_ready_to_write_dly = is_lst_kk_dly[P_ICH] && is_lst_fi_dly[P_ICH];

    // Pipeline enable: can accept new data when input is valid or processing
    assign pipeline_enable           = (is_fst_fo ? in_valid : 1'b1) && (!output_ready_to_write_dly || out_ready);

    // Input ready when we can accept data
    assign in_ready                  = pipeline_enable && is_fst_fo;

    // Weight address calculation
    assign weight_addr               = (cntr_fo * FOLD_I * KK) + (cntr_fi * KK) + cntr_kk;

    // Main pipeline process
    always_ff @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            cntr_hw <= 0;
            cntr_fo <= 0;
            cntr_fi <= 0;
            cntr_kk <= 0;

            // Clear line buffer
            for (int i = 0; i < FOLD_I; i++) begin
                for (int k = 0; k < KK; k++) begin
                    line_buffer[i][k] <= '0;
                end
            end

        end else if (pipeline_enable) begin

            // ============== Load input data ==============
            if (is_fst_fo) begin
                line_buffer[cntr_fi][cntr_kk] <= in_data;
            end

            // ============== Counter updates ==============
            if (cntr_kk == KK - 1) begin
                cntr_kk <= 0;

                if (cntr_fi == FOLD_I - 1) begin
                    cntr_fi <= 0;

                    if (cntr_fo == FOLD_O - 1) begin
                        cntr_fo <= 0;

                        if (cntr_hw == N_HW - 1) begin
                            cntr_hw <= 0;
                        end else begin
                            cntr_hw <= cntr_hw + 1;
                        end
                    end else begin
                        cntr_fo <= cntr_fo + 1;
                    end
                end else begin
                    cntr_fi <= cntr_fi + 1;
                end
            end else begin
                cntr_kk <= cntr_kk + 1;
            end
        end
    end

    // Delay registers for input data and control signals to match weight_data latency
    logic [P_ICH*A_BIT-1:0] in_buf_d;
    logic                   is_fst_kk_d;
    logic                   pipeline_enable_d;

    always_ff @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            in_buf_d          <= '0;
            is_fst_kk_d       <= 1'b0;
            pipeline_enable_d <= 1'b0;
        end else begin
            in_buf_d          <= is_fst_fo ? in_data : line_buffer[cntr_fi][cntr_kk];
            is_fst_kk_d       <= (cntr_kk == 0) && (cntr_fi == 0);
            pipeline_enable_d <= pipeline_enable;
        end
    end

    // MAC array instances for each output channel
    logic        [A_BIT-1:0] x_vec[P_ICH];
    logic signed [W_BIT-1:0] w_vec[P_OCH] [P_ICH];

    // Unpack delayed input data to array
    always_comb begin
        for (int i = 0; i < P_ICH; i++) begin
            x_vec[i] = in_buf_d[i*A_BIT+:A_BIT];
        end
    end

    // Unpack weight data to array
    always_comb begin
        for (int o = 0; o < P_OCH; o++) begin
            for (int i = 0; i < P_ICH; i++) begin
                w_vec[o][i] = weight_data[(P_ICH*o+i)*W_BIT+:W_BIT];
            end
        end
    end

    // Instantiate MAC arrays for each output channel
    generate
        for (genvar o = 0; o < P_OCH; o++) begin : gen_mac_array
            mac_array #(
                .P_ICH(P_ICH),
                .A_BIT(A_BIT),
                .W_BIT(W_BIT),
                .B_BIT(B_BIT)
            ) u_mac_array (
                .clk  (clk),
                .rst_n(rst_n),
                .en   (pipeline_enable_d),
                .clr  (is_fst_kk_d),
                .x_vec(x_vec),
                .w_vec(w_vec[o]),
                .acc  (acc[o])
            );
        end
    endgenerate

    // Output generation
    always_ff @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            out_data  <= '0;
            out_valid <= 1'b0;
        end else begin
            if (output_ready_to_write_dly && (!out_valid || out_ready)) begin
                // Pack accumulator outputs
                for (int o = 0; o < P_OCH; o++) begin
                    out_data[o*B_BIT+:B_BIT] <= acc[o];
                end
                out_valid <= 1'b1;
            end else if (out_ready) begin
                out_valid <= 1'b0;
            end
        end
    end

endmodule
